package developer.zyk.ability.knowledge.parsers.textParser;

import lombok.Data;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.document.Document;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.regex.Pattern;

/**
 * @author zhangyongkang
 * @time 2025/5/20 11:02
 * @description
 */
@Component
public class TextParser {

    public List<Document> parse(String content) {
        // 创建要分割的标题列表
        List<AbstractMap.SimpleEntry<String, String>> headers = new ArrayList<>();
        headers.add(new AbstractMap.SimpleEntry<>("#", "header_1"));
        headers.add(new AbstractMap.SimpleEntry<>("##", "header_2"));
        headers.add(new AbstractMap.SimpleEntry<>("###", "header_3"));


        MarkdownHeaderTextSplitter markdownHeaderTextSplitter = new MarkdownHeaderTextSplitter(headers);
        return markdownHeaderTextSplitter.splitText(content);
    }


    /**
     * 根据指定的标题分割Markdown文件
     *
     * @author 粒方
     * @since 2025/3/9
     **/
    public class MarkdownHeaderTextSplitter {
        private static List<AbstractMap.SimpleEntry<String, String>> headersToSplitOn;
        private boolean returnEachLine;
        private boolean stripHeaders;
        private static final Pattern CONTROL_CHARS =
                Pattern.compile("[\\x00-\\x1F\\x7F\\p{Cc}]");


        /**
         * 使用指定参数创建新的MarkdownHeaderTextSplitter实例
         *
         * @param headersToSplitOn 需要追踪的标题
         * @param returnEachLine   是否返回每行及其关联的标题
         * @param stripHeaders     是否从块内容中删除分割标题
         */
        public MarkdownHeaderTextSplitter(
                List<AbstractMap.SimpleEntry<String, String>> headersToSplitOn,
                boolean returnEachLine,
                boolean stripHeaders) {
            // 按长度降序排序标题
            headersToSplitOn.sort((a, b) -> Integer.compare(b.getKey().length(), a.getKey().length()));

            this.headersToSplitOn = headersToSplitOn;
            this.returnEachLine = returnEachLine;
            this.stripHeaders = stripHeaders;
        }

        /**
         * 使用指定标题创建新的MarkdownHeaderTextSplitter实例
         *
         * @param headersToSplitOn 需要追踪的标题
         */
        public MarkdownHeaderTextSplitter(List<AbstractMap.SimpleEntry<String, String>> headersToSplitOn) {
            this(headersToSplitOn, false, true);
        }

        /**
         * 表示带有关联元数据的文本行
         */
        @Data
        private static class LineType {
            private String content;
            private Map<String, Object> metadata;

            public LineType(String content, Map<String, Object> metadata) {
                this.content = content;
                this.metadata = metadata;
            }

        }

        /**
         * 表示Markdown文档中的标题
         */
        @Data
        private static class HeaderType {
            private int level;
            private String name;
            private String data;

            public HeaderType(int level, String name, String data) {
                this.level = level;
                this.name = name;
                this.data = data;
            }

        }

        /**
         * 将具有相同元数据的行组合成块
         *
         * @param lines 带有标题元数据的文本行
         * @return Document对象列表
         */
        private List<Document> aggregateLinesToChunks(List<LineType> lines) {
            List<LineType> aggregatedChunks = new ArrayList<>();

            for (LineType line : lines) {
                if (!aggregatedChunks.isEmpty()) {

                    LineType lastChunk = aggregatedChunks.get(aggregatedChunks.size() - 1);
                    Map<String, Object> lastMetadata = lastChunk.getMetadata();
                    Map<String, Object> currentMetadata = line.getMetadata();

                    if (lastMetadata.equals(currentMetadata)) {
                        // 如果聚合列表中的最后一行与当前行具有相同的元数据，将当前内容附加到最后一行的内容中
                        lastChunk.setContent(lastChunk.getContent() + "  \n" + line.getContent());
                    } else if (lastMetadata.size() < currentMetadata.size()
                            && lastChunk.getContent().split("\n")[lastChunk.getContent().split("\n").length - 1].startsWith("#")
                            && !stripHeaders) {
                        // 如果聚合列表中的最后一行具有与当前行不同的元数据，并且具有比当前行浅的标题级别，最后一行是一个标题，我们也不会剥去标头，将当前内容附加到最后一行的内容
                        lastChunk.setContent(lastChunk.getContent() + "  \n" + line.getContent());
                        // 并更新最后一行的元数据
                        lastChunk.setMetadata(currentMetadata);
                    } else {
                        // 否则，将当前行添加到聚合列表中
                        aggregatedChunks.add(line);
                    }
                } else {
                    // 如果聚合列表为空，将当前行添加到聚合列表中
                    aggregatedChunks.add(line);
                }
            }

            // 将聚合的块转换为Document对象列表
            return aggregatedChunks.stream()
                    .map(chunk -> new Document(chunk.getContent(), chunk.getMetadata()))
                    .toList();
        }

        /**
         * 基于标题将Markdown文件分割成块
         *
         * @param text Markdown文件内容
         * @return Document对象列表
         */
        public List<Document> splitText(String text) {
            // 按行分割文本
            String[] lines = text.split("\n");
            // 分块结果
            List<LineType> linesWithMetadata = new ArrayList<>();
            // 当前正在处理的数据块的内容和元数据
            List<String> currentContent = new ArrayList<>();
            Map<String, String> currentMetadata = new HashMap<>();
            // 标题栈
            List<HeaderType> headerStack = new ArrayList<>();
            Map<String, String> initialMetadata = new HashMap<>();

            // 是否编码块
            boolean inCodeBlock = false;
            String openingFence = "";

            for (String line : lines) {
                // 处理原文件格式，去除不可见字符
                String strippedLine = line.strip();
                strippedLine = CONTROL_CHARS.matcher(strippedLine).replaceAll("");

                // 处理代码块
                if (!inCodeBlock) {
                    if (strippedLine.startsWith("```") && StringUtils.countMatches(line, "```") == 1) {
                        inCodeBlock = true;
                        openingFence = "```";
                    } else if (strippedLine.startsWith("~~~")) {
                        inCodeBlock = true;
                        openingFence = "~~~";
                    }
                } else {
                    if (strippedLine.startsWith(openingFence)) {
                        inCodeBlock = false;
                        openingFence = "";
                    }
                }

                if (inCodeBlock) {
                    currentContent.add(strippedLine);
                    continue;
                }

                // 根据每个标题类型检查每一行（例如，#，##）
                boolean foundHeader = false;
                for (AbstractMap.SimpleEntry<String, String> headerToSplitOn : headersToSplitOn) {
                    String sep = headerToSplitOn.getKey();
                    String name = headerToSplitOn.getValue();

                    // 判断是否标题行
                    if (strippedLine.startsWith(sep) && (
                            // 1、# 没有标题文字
                            // 2、标题行后面是空格
                            strippedLine.length() == sep.length() || strippedLine.charAt(sep.length()) == ' '
                    )) {
                        if (StringUtils.isNotEmpty(name)) {
                            // 获取当前行标题级别
                            int currentHeaderLevel = StringUtils.countMatches(sep, "#");

                            while (!headerStack.isEmpty() &&
                                    headerStack.get(headerStack.size() - 1).getLevel() >= currentHeaderLevel) {
                                // 从标题栈中移除最后一个标题

                                HeaderType poppedHeader = headerStack.remove(headerStack.size()-1);
                                initialMetadata.remove(poppedHeader.getName());
                            }
                            String headerText = strippedLine.substring(sep.length()).trim();
                            headerStack.add(new HeaderType(currentHeaderLevel, name, headerText));
                            initialMetadata.put(name, headerText);
                        }

                        if (!currentContent.isEmpty()) {
                            linesWithMetadata.add(new LineType(String.join("\n", currentContent), new HashMap<>(currentMetadata)));
                            currentContent.clear();
                        }

                        if (!stripHeaders) {
                            currentContent.add(strippedLine);
                        }
                        foundHeader = true;
                        break;
                    }
                }

                if (!foundHeader) {
                    if (StringUtils.isNotEmpty(strippedLine)) {
                        currentContent.add(strippedLine);
                    } else if (!currentContent.isEmpty()) {
                        linesWithMetadata.add(new LineType(String.join("\n", currentContent), new HashMap<>(currentMetadata)));
                        currentContent.clear();
                    }

                }

                currentMetadata = new HashMap<>(initialMetadata);
            }

            if (!currentContent.isEmpty()) {
                linesWithMetadata.add(new LineType(String.join("\n", currentContent), new HashMap<>(currentMetadata)));
            }

            if (!returnEachLine) {
                return aggregateLinesToChunks(linesWithMetadata);
            } else {
                return linesWithMetadata.stream()
                        .map(line -> new Document(line.getContent(), line.getMetadata()))
                        .toList();
            }
        }
    }

}
